package org.seedframework.sharding.mybatis;

import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.executor.SimpleExecutor;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.transaction.jdbc.JdbcTransaction;
import org.mybatis.spring.transaction.SpringManagedTransaction;
import org.seedframework.sharding.resource.DoubleWriteDataSourcePoolContext;
import org.seedframework.sharding.resource.ShardingDataSource;
import org.seedframework.sharding.resource.ShardingDataSourcePoolContext;
import org.springframework.context.ApplicationContext;
import org.springframework.jdbc.datasource.DataSourceTransactionManager;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import java.sql.SQLException;
import java.util.*;

/**
 * Set ibatis Executor data source
 */
@Slf4j
public class ExecutorFactory {

    private static Map<String, Executor> executorMap = new WeakHashMap<>();

    private static Map<String, Set<Executor>> executorSetMap = new WeakHashMap<>();

    private static Configuration configuration = null;

    private static ApplicationContext applicationContext;

    private static ThreadLocal<DataSourceTransactionManager> currentDataSourceTransactionManager = new ThreadLocal<>();

    public static Executor getExecutor(String id, Executor executor) throws SQLException {

        Map<String, ShardingDataSource> dataSourceMap = ShardingDataSourcePoolContext.getDataSourceWithKey();
        if (CollectionUtils.isEmpty(dataSourceMap)) {
            return executor;
        }
        ShardingDataSource dataSource = dataSourceMap.get(id);
        if (dataSource == null) {
            return executor;
        }
        Assert.notNull(configuration, "MyBatis configuration is not found in application-context.");
        DataSourceTransactionManager transactionManager = new DataSourceTransactionManager();
        transactionManager.setDataSource(dataSource);
        currentDataSourceTransactionManager.set(transactionManager);
        // spring-framework @Transactional annotation support
        Executor cached = executorMap.get(id);
        if (cached != null) {
            return cached;
        }


        cached = new SimpleExecutor(configuration, new SpringManagedTransaction(dataSource));
        executorMap.put(id, cached);
        return cached;

    }

    public static Set<Executor> getExecutors(String groupName, Executor executorDefault) throws SQLException {
        Set<Executor> rst = new HashSet<>();
        Set<ShardingDataSource> set = DoubleWriteDataSourcePoolContext.getDataSourceSet(groupName);
        Assert.notEmpty(set, "Double-write data source is empty.");
        Set<Executor> cached = executorSetMap.get(groupName);
        if (!CollectionUtils.isEmpty(cached)) {
            return cached;
        }
        for (ShardingDataSource dataSource : set) {
            Executor executor = new SimpleExecutor(configuration, new SpringManagedTransaction(dataSource));
            rst.add(executor);
        }
        if (rst.isEmpty()) {
            rst.add(executorDefault);
        }
        executorSetMap.put(groupName, rst);
        return rst;
    }


    public static void setConfiguration(Configuration config) {
        ExecutorFactory.configuration = config;
    }


    public static void setApplicationContext(ApplicationContext applicationContext){
        ExecutorFactory.applicationContext = applicationContext;
    }

    public static ApplicationContext getApplicationContext(){
        return applicationContext;
    }

    public static DataSourceTransactionManager getCurrentTransactionManager(){
        return currentDataSourceTransactionManager.get();
    }


}
